{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ba0b121f",
   "metadata": {},
   "source": [
    "## Dendrite Segmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08845ee3",
   "metadata": {},
   "source": [
    "### Image Binarization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06817dc1",
   "metadata": {},
   "source": [
    "Set image sampling density in micrometer/pixel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a6930df",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampling_x = 0.025\n",
    "sampling_y = 0.025\n",
    "sampling_z = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efdd9c4c",
   "metadata": {},
   "source": [
    "Load dendrite image file.<br>\n",
    "Pass .tif file path into *load_tif* function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3395c55",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from spine_segmentation import load_tif\n",
    "from notebook_widgets import show_3d_image\n",
    "from scipy.ndimage import gaussian_filter\n",
    "import os\n",
    "\n",
    "image_path = \"example_dendrite/example_dendrite.tif\"\n",
    "image_name = os.path.basename(image_path)\n",
    "\n",
    "# load tif\n",
    "image: np.ndarray = load_tif(image_path)\n",
    "image = gaussian_filter(image, 1)\n",
    "show_3d_image(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42acfc5f",
   "metadata": {},
   "source": [
    "Perform image binarization.<br>\n",
    "Parameters:<br>\n",
    "* $BaseThreshold$ — base threshold value\n",
    "* $Weight$ — how much neighbouring values affect threshold\n",
    "* $BlockSize$ — size of the neighbourhood area to calculate median value in\n",
    "\n",
    "Local binarization is calculated as follows:<br>\n",
    "\n",
    "$LocalThreshold_{xyz} = BaseThreshold + Weight \\cdot (BaseThreshold - LocalMedianValue_{xyz}(BlockSize))$<br>\n",
    "$BinarizedValue_{xyz} = \\begin{cases} 1\\text{, }Value_{xyz} > LocalThreshold_{xyz} \\\\ 0\\text{, else} \\end{cases}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4894f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from notebook_widgets import interactive_binarization\n",
    "\n",
    "binarization_widget = interactive_binarization(image)\n",
    "display(binarization_widget)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94876cb8",
   "metadata": {},
   "source": [
    "Select connected component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "844404d0",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from notebook_widgets import select_connected_component_widget\n",
    "\n",
    "select_component_widget = select_connected_component_widget(binarization_widget.result)\n",
    "display(select_component_widget)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3faf57a9",
   "metadata": {},
   "source": [
    "### Construct 3D surface to perform segmention on.\n",
    "\n",
    "First, from binarized image, calculate points that belong to the surface, as well as surface normals in those points.\n",
    "Use Poisson surface reconstruction algorithm to generate the surface mesh.<br>\n",
    "Algorithm takes set of points with normals and produces a smooth closed surface mesh $S$.<br>\n",
    "Generated mesh is saved to <i>\"output/surface_mesh.off\"</i> file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4a7137d",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from spine_segmentation import get_surface_points\n",
    "from CGAL.CGAL_Point_set_3 import Point_set_3\n",
    "from CGAL.CGAL_Polyhedron_3 import Polyhedron_3\n",
    "from CGAL.CGAL_Poisson_surface_reconstruction import poisson_surface_reconstruction\n",
    "from notebook_widgets import show_3d_mesh, create_dir\n",
    "import os\n",
    "from spine_segmentation import save_tif, apply_scale \n",
    "\n",
    "\n",
    "# extract binarization result\n",
    "binary_image = select_component_widget.result\n",
    "\n",
    "# export surface mesh to .off file\n",
    "save_path = f\"output/{image_name.rsplit('.', 1)[0]}\"\n",
    "\n",
    "create_dir(\"output\")\n",
    "create_dir(save_path)\n",
    "\n",
    "save_tif(f\"{save_path}/{image_name}\", image)\n",
    "save_tif(f\"{save_path}/binarized_image.tif\", binary_image)\n",
    "\n",
    "z_display_factor = 0.5\n",
    "\n",
    "# calculate surface points\n",
    "# (only apply sampling scale after surface reconstruction, because small coordinates give worse reconstruction results smh)\n",
    "# surface_points: Point_set_3 = get_surface_points(binary_image, (sampling_x, sampling_y, sampling_z))\n",
    "surface_points: Point_set_3 = get_surface_points(binary_image)\n",
    "\n",
    "# construct surface mesh\n",
    "surface_poly = Polyhedron_3()\n",
    "poisson_surface_reconstruction(surface_points, surface_poly)\n",
    "\n",
    "surface_poly = apply_scale(surface_poly, (sampling_x, sampling_y, sampling_z))\n",
    "surface_poly.write_to_file(f\"{save_path}/surface_mesh.off\")\n",
    "\n",
    "# render surface mesh\n",
    "show_3d_mesh(apply_scale(surface_poly, (1, 1, 1 * z_display_factor)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9569f82",
   "metadata": {},
   "source": [
    "### Mesh segmentation.\n",
    "\n",
    "Parameters:\n",
    "* Sensitivity — how much distance from skeleton affects segmentation. Higher values will result in less false positive spines, but worse segmentation at spine base and detection of smaller spines. \n",
    "\n",
    "\n",
    "Algorith is as follows:\n",
    "1. Calculate mesh skeleton. Mean Curvature Skeleton algorithm is used. Algorithm generates skeleton graph $G$ and correspondence from each point on the surface to some point in the skeleton $f:S\\rightarrow G$.\n",
    "2. Find dendrite skeleton subgraph $G_{dendrite}$ — longest path in the graph, with the least sum angle between consecutive edges.\n",
    "3. Mark surface points that don't correspond to dendrite skeleton subgraph as spines. $S_{spines} = \\{ p \\mid p \\in S \\land f(p) \\notin G_{dendrite} \\}$\n",
    "4. Calculate distance from surface to skeleton statistic. $Dist = \\{ \\| p-f(p) \\| \\mid p \\in S \\}$\n",
    "5. Mark surface points that are futher away from skeleton than others as spines. $S_{spines} = S_{spines} \\cup \\{ p \\mid p \\in S \\land \\| p - f(p) \\| > quantile(Dist, Sensitivity) \\}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27ea84ee",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from CGAL.CGAL_Surface_mesh_skeletonization import surface_mesh_skeletonization\n",
    "from CGAL.CGAL_Polygon_mesh_processing import Polylines\n",
    "from spine_segmentation import build_graph, build_correspondence, build_reverse_correpondnce\n",
    "from notebook_widgets import interactive_segmentation\n",
    "\n",
    "# reload mesh because saving as .off compresses point coordinates\n",
    "# and we want exactly same coordinates for spines and denrite\n",
    "surface_poly = Polyhedron_3(f\"{save_path}/surface_mesh.off\")\n",
    "surface_poly = apply_scale(surface_poly, (1, 1, 1 * z_display_factor))\n",
    "\n",
    "# get surface skeleton\n",
    "skeleton_polylines = Polylines()\n",
    "correspondence_polylines = Polylines()\n",
    "surface_mesh_skeletonization(surface_poly, skeleton_polylines, correspondence_polylines)\n",
    "\n",
    "# convert to more performant data format \n",
    "skeleton_graph = build_graph(skeleton_polylines)\n",
    "corr = build_correspondence(correspondence_polylines)\n",
    "reverse_corr = build_reverse_correpondnce(correspondence_polylines)\n",
    "\n",
    "# perform segmentation \n",
    "segmentation_widget = interactive_segmentation(surface_poly, corr, reverse_corr, skeleton_graph)\n",
    "display(segmentation_widget)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9985e30d",
   "metadata": {},
   "source": [
    "Generate surface mesh for each individual spine.<br>\n",
    "Calculate metrics for each spine. Calculated metric names are defined in the *metric_names* list.<br>\n",
    "Manually remove false positive spine selections.<br>\n",
    "Use **index** to navigate spines, **checkbox** to keep or remove spine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e29cdb",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from spine_segmentation import get_spine_meshes\n",
    "from spine_metrics import calculate_metrics\n",
    "from notebook_widgets import select_spines_widget\n",
    "\n",
    "# extract segmentation result\n",
    "segmentation = segmentation_widget.result\n",
    "\n",
    "# extract spine meshes\n",
    "spine_meshes = get_spine_meshes(surface_poly, segmentation)\n",
    "\n",
    "# define calculated metrics\n",
    "metric_names = [\"Area\", \"Volume\"]\n",
    "\n",
    "# manually select valid spines\n",
    "selection_widget = select_spines_widget(spine_meshes, surface_poly, metric_names)\n",
    "display(selection_widget)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae4e988f",
   "metadata": {},
   "source": [
    "Generate final segmentation from manually filtered spines. Spine meshes are saved to <i>\"output/spine_{i}.off\"</i> files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "254f9a2a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from spine_segmentation import spines_to_segmentation, save_segmentation \n",
    "from notebook_widgets import show_segmented_mesh\n",
    "\n",
    "\n",
    "# extract selected spines\n",
    "selected_spines = []\n",
    "\n",
    "# export selected spine meshes to .off files\n",
    "for i, (spine_mesh, spine_metrics) in enumerate(selection_widget.children[1].result):\n",
    "    filename = f\"{save_path}/spine_{i}.off\"\n",
    "    scaled_spine_mesh = apply_scale(spine_mesh, (1, 1, 1 / z_display_factor))  \n",
    "    scaled_spine_mesh.write_to_file(filename)\n",
    "    selected_spines.append(spine_mesh)\n",
    "\n",
    "# combine selected spines into segmentation\n",
    "manually_filtered_segmentation = spines_to_segmentation(selected_spines)\n",
    "save_segmentation(manually_filtered_segmentation, f\"{save_path}/segmentation.json\")\n",
    "\n",
    "# render manually filtered segmentation\n",
    "show_segmented_mesh(surface_poly, manually_filtered_segmentation)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
